#!/usr/bin/python2
import sys, base64, struct, xml.parsers.expat

from core import MSScan

class MzXML:
    def __init__(self):
        self.msLevel = 0
        self.current_tag = ''
        self.tag_level = 0
        self._precursors = {}
        self.MS_list = []

    def decode_spectrum(self, line):
        decoded = base64.decodestring(line)
        tmp_size = len(decoded)/4
        unpack_format1 = ">%dL" % tmp_size

        idx = 0
        mz_list = []
        intensity_list = []

        for tmp in struct.unpack(unpack_format1,decoded):
            tmp_i = struct.pack("I", tmp)
            tmp_f = struct.unpack("f", tmp_i)[0]
            if( idx % 2 == 0 ):
                mz_list.append( float(tmp_f) )
            else:
                intensity_list.append( float(tmp_f) )
            idx += 1

        return mz_list, intensity_list

    def _start_element(self, name, attrs):
        self.tag_level += 1
        self.current_tag = name
        if name == 'precursorMz':
            self.MS_list[-1].precursor_intensity = float(attrs['precursorIntensity'])

        if name == 'scan':
            self.msLevel = int(attrs['msLevel'])
            tmp_ms = MSScan()

            tmp_ms.ms_level = self.msLevel
            tmp_ms.polarity = attrs.setdefault('polarity', '+')
            tmp_ms.scan_num = int(attrs['num'])
            tmp_ms.peak_count = int(attrs['peaksCount'])
            tmp_ms.filter_line = attrs.setdefault('filterLine', '')
            tmp_ms.retention_time = float(attrs['retentionTime'].strip('PTS'))
            tmp_ms.low_mz = float(attrs['lowMz'])
            tmp_ms.high_mz = float(attrs['highMz'])
            tmp_ms.base_peak_mz = float(attrs['basePeakMz'])
            tmp_ms.base_peak_intensity = float(attrs['basePeakIntensity'])
            tmp_ms.total_ion_current = float(attrs['totIonCurrent'])
            tmp_ms.encoded_mz = ''
            tmp_ms.encoded_intensity = ''
            self._precursors[self.msLevel] = tmp_ms.scan_num
            if self.msLevel != 1:
                tmp_ms.collision_energy = float(attrs['collisionEnergy'])
                tmp_ms.precursor_scan_num = self._precursors.setdefault(self.msLevel-1, 0)

            self.MS_list.append(tmp_ms)

    def _end_element(self, name):
        self.tag_level -= 1
        self.current_tag = ''
        self.msLevel == 0

    def _char_data(self, data):
        if self.current_tag == 'precursorMz':
            self.MS_list[-1]._precursor_mz0 = float(data)

        if self.current_tag == 'peaks':
            mz_list, intensity_list = self.decode_spectrum(data)
            mz_string = ''.join([struct.pack('>f', i) for i in mz_list])
            intensity_string = ''.join([struct.pack('>f', i) for i in intensity_list])
            self.MS_list[-1].encoded_mz += mz_string
            self.MS_list[-1].encoded_intensity += intensity_string

    def parse_file(self, filename_xml):
        sys.stderr.write("Read %s ... "%filename_xml)
        f_xml = open(filename_xml,'r')
        content_xml = ''
        for line in f_xml:
            content_xml += line
        f_xml.close()
        sys.stderr.write("Done\n")

        expat = xml.parsers.expat.ParserCreate()
        expat.StartElementHandler = self._start_element
        expat.EndElementHandler = self._end_element
        expat.CharacterDataHandler = self._char_data
        expat.Parse(content_xml)
